/*
 * Generate Qingtian QTSM Attestation Doc
 *
 * Copyright (c) Huawei Technologies Co., Ltd. 2022. All Rights Reserved.
 * SPDX-License-Identifier: Apache-2.0.
 */
#include <openssl/bn.h>
#include <openssl/evp.h>
#include <openssl/rsa.h>
#include <openssl/pem.h>
#include <string.h>
#include <stdbool.h>

#include "attestation.h"
#include "qtsm_lib.h"

static bool stream_buf_is_valid(const struct stream_buf *binary_in)
{
    if (binary_in == NULL ||
       !((binary_in->capacity == 0 && binary_in->len == 0 && binary_in->buffer == NULL) ||
        (binary_in->capacity > 0 && binary_in->len <= binary_in->capacity &&
        binary_in->buffer != NULL)))
        return false;

    return true;
}

static bool rsa_keypair_is_valid(const struct rsa_keypair *keypair)
{
    if (keypair == NULL || keypair->keys == NULL)
        return false;
    return true;
}

struct rsa_keypair *attestation_rsa_keypair_new(enum rsa_key_size key_size)
{
    RSA *rsa = NULL;
    BIGNUM *e = NULL;
    EVP_PKEY *keypair = NULL;
    struct rsa_keypair *key = NULL;

    rsa = RSA_new();
    if (rsa == NULL)
        return NULL;

    e = BN_new();
    if (e == NULL) {
        RSA_free(rsa);
        return NULL;
    }
    if (!BN_set_word(e, RSA_F4)) {
        RSA_free(rsa);
        return NULL;
    }

    if (RSA_generate_key_ex(rsa, key_size, e, NULL) != 1) {
        BN_free(e);
        RSA_free(rsa);
        return NULL;
    }

    BN_free(e);

    keypair = EVP_PKEY_new();
    if (keypair == NULL || EVP_PKEY_assign_RSA(keypair, rsa) != 1) {
        RSA_free(rsa);
        return NULL;
    }

    key = malloc(sizeof(struct rsa_keypair));
    if (!key) {
        EVP_PKEY_free(keypair);
        return NULL;
    }
    memset(key, 0, sizeof(struct rsa_keypair));

    key->keys = keypair;

    return key;
}

void attestation_rsa_keypair_destroy(struct rsa_keypair *keypair)
{
    if (!rsa_keypair_is_valid(keypair))
        return;

    EVP_PKEY_free(keypair->keys);
    free(keypair);
}

int get_attestation_doc(const struct rsa_keypair *keypair,
                        struct stream_buf *attestation_doc)
{
    int qtsm_dev_fd;
    int rc;
    u_char *user_pubkey = NULL;
    u_int user_pubkey_len;
    BIO *temp_buf = NULL;

    if (!rsa_keypair_is_valid(keypair) ||
        !stream_buf_is_valid(attestation_doc))
        return CMS_INVALID_ARGUMENT;

    temp_buf = BIO_new(BIO_s_mem());
    if (!temp_buf)
        return INTERNAL_ERROR;

    rc = i2d_PUBKEY_bio(temp_buf, keypair->keys);
    if (!rc) {
        rc = INTERNAL_ERROR;
        goto err;
    }

    user_pubkey_len = BIO_get_mem_data(temp_buf, &user_pubkey);
    if (user_pubkey_len <= 0) {
        rc = INTERNAL_ERROR;
        goto err;
    }

    qtsm_dev_fd = qtsm_lib_init();
    if (qtsm_dev_fd < 0) {
        rc = INTERNAL_ERROR;
        goto err;
    }

    attestation_doc->len = attestation_doc->capacity;
    rc = qtsm_get_attestation(qtsm_dev_fd, NULL, 0, NULL, 0,
        (const u_char *)user_pubkey, user_pubkey_len,
        attestation_doc->buffer, (uint32_t *)&attestation_doc->len);

    qtsm_lib_exit(qtsm_dev_fd);

err:
    BIO_free(temp_buf);
    return rc;
}

int attestation_rsa_decrypt(struct rsa_keypair *keypair,
    const struct stream_buf *cipher_symm_key,
    struct stream_buf *plain_symm_key)
{
    size_t ek_len;
    unsigned char *ek = NULL;
    EVP_PKEY *key = NULL;
    EVP_PKEY_CTX *pkey_ctx = NULL;
    int rc;

    if (!rsa_keypair_is_valid(keypair) ||
        !stream_buf_is_valid(cipher_symm_key))
        return CMS_INVALID_ARGUMENT;

    key = keypair->keys;

    /* Create the decryption context */
    pkey_ctx = EVP_PKEY_CTX_new(key, NULL);
    if (!pkey_ctx)
        return CMS_INTERNAL_ERR;

    /* Initialize the context */
    if (EVP_PKEY_decrypt_init(pkey_ctx) != 1 ||
        EVP_PKEY_CTX_set_rsa_padding(pkey_ctx, RSA_PKCS1_OAEP_PADDING) != 1 ||
        EVP_PKEY_CTX_set_rsa_mgf1_md(pkey_ctx, EVP_sha256()) != 1 ||
        EVP_PKEY_CTX_set_rsa_oaep_md(pkey_ctx, EVP_sha256()) != 1) {
        rc = CMS_INTERNAL_ERR;
        goto err;
    }

    /* Decrypt */
    if (EVP_PKEY_decrypt(pkey_ctx, NULL, &ek_len,
        cipher_symm_key->buffer, cipher_symm_key->len) <= 0) {
        rc = CMS_INTERNAL_ERR;
        goto err;
    }

    ek = OPENSSL_malloc(ek_len);
    if (!ek) {
        rc = CMS_INTERNAL_ERR;
        goto err;
    }

    if (EVP_PKEY_decrypt(pkey_ctx, ek, &ek_len,
        cipher_symm_key->buffer, cipher_symm_key->len) <= 0) {
        rc = CMS_INTERNAL_ERR;
        goto err;
    }

    plain_symm_key->capacity = plain_symm_key->len = ek_len;
    plain_symm_key->buffer = malloc(plain_symm_key->capacity);
    if (!plain_symm_key->buffer) {
        rc = CMS_INTERNAL_ERR;
        goto err;
    }
    memcpy(plain_symm_key->buffer, ek, ek_len);

    rc = CMS_SUCCESS;

err:
    EVP_PKEY_CTX_free(pkey_ctx);
    if (ek != NULL)
        OPENSSL_free(ek);

    return rc;
}
